# koroFileHeader at Yungoal acer
# Create: 2021-05-28 09:42:07
# LastEdit: 2025-11-03 14:43:03
"""使用 Azure 的 run command 检查进程
用 ps 命令检查linux系统的进程
用 systemctl is-enabled 或 chkconfig 检查linux系统的服务
用 Get-Process 检查windows系统的进程
用 Get-Service 检查windows系统的服务状态
"""
__author__ = '749B'

import os, sys
BASE_DIR = os.getcwd()
sys.path.append(BASE_DIR)

import argparse
import textwrap
import msrestazure.azure_cloud
# pylint: disable = import-error
from azure_client import Azure_Client
from azure_client import data_format
from azure_client.models import RuncmdAsyncJob
from utils.tools import remove_empty
from utils.blueking import job_start, job_success, job_fail
from utils.env import init_environs

from typing import Optional, List
# 只是因为函数的返回需要做类型判断
from azure.mgmt.resourcegraph.models._models_py3 import QueryResponse


def str_query_base(
    is_linux: bool = True,
    subscription_id: Optional[str] = None,
    resource_group: Optional[str] = None,
    filters: Optional[str] = None,
    ) -> str:
    """包含过滤条件的查询语句，在这个基础上过滤虚拟机"""
    strQuery = textwrap.dedent("""\
    Resources
    | where type == "microsoft.compute/virtualmachines"
    | extend 
        OSType = properties.storageProfile.osDisk.osType,
        osdiskid = tolower(tostring(properties.storageProfile.osDisk.managedDisk.id)),
        osdiskvhduri = properties.storageProfile.osDisk.vhd.uri,
        nicid = tostring(properties.networkProfile.networkInterfaces[0].id)
    | where OSType =~ "{os_type}"
    {filter_subscription_id}
    {filter_resource_group}
    {other_filters}
    | project vm_name=name, resourceGroup, OSType, subscriptionId, nicid, osdiskid, osdiskvhduri
    | join kind=leftouter (
        Resourcecontainers
        | where type =~ "microsoft.resources/subscriptions"
        | project SubName=name, subscriptionId
    ) on subscriptionId
    | project-away subscriptionId1
    | join kind=leftouter (
        Resources
        | where type =~ 'microsoft.network/networkinterfaces'
        | project nicid=id, ip_address=properties.ipConfigurations[0].properties.privateIPAddress
    ) on nicid 
    | project-away nicid, nicid1
    | join kind=leftouter (
        Resources
        | where type == "microsoft.compute/disks"
        | project osdiskid=tolower(id), osdisk_creattime=properties.timeCreated
    ) on osdiskid
    | project-away osdiskid, osdiskid1
    """).format(
        os_type='linux' if is_linux else 'windows',
        filter_subscription_id=f'| where subscriptionId =~ "{subscription_id}"' if subscription_id else '',
        filter_resource_group=f'| where resourceGroup =~ "{resource_group}"' if resource_group else '',
        other_filters=f'| where {filters}' if filters else ''
        )
    # 去掉空白行
    strQuery = remove_empty(strQuery)
    return strQuery
    
def get_vm_count(client: Azure_Client, base_query: str) -> int:
    """返回虚拟机的数量"""
    strQuery = base_query + '| summarize count()\n'
    # 去掉空白行
    strQuery = remove_empty(strQuery)
    res = client.resource_graph_query_default(
        strQuery,
        [sub.get("subscription_id") for sub in client.get_subscriptions()],
    )
    return res.data[0].get('count_')

def get_vm_info(
    client: Azure_Client, 
    base_query: str,
    limit: Optional[int] = None,
    ) -> QueryResponse:
    """拉取所有的虚拟机"""
    strQuery = base_query + f'| limit {limit}\n' if limit else base_query 
    # 去掉空白行
    strQuery = remove_empty(strQuery)
    res = client.resource_graph_query_default(
        strQuery,
        [sub.get("subscription_id") for sub in client.get_subscriptions()],
    )
    return res

def init_argparse(args: Optional[List[str]] = None) -> argparse.Namespace:
    """初始化：命令行参数"""
    parser = argparse.ArgumentParser(description='检查服务/进程')

    parser.add_argument('services', nargs='+', help='服务或进程的名称')
    parser.add_argument('-e', '--envfile', help="存有用户认证信息的env文件的文件名")
    parser.add_argument('-s', '--subscription', help='限定订阅，订阅名称或订阅ID都可')
    parser.add_argument('-g', '--resource_group', help='限定资源组名称')
    parser.add_argument('-f', '--filters', help='其他过滤条件，参考资源查询的where语句')
    parser.add_argument('-c', '--command', default='ps', choices=list(SCRIPTS.keys()), help='使用的检查命令')
    parser.add_argument('-t', '--timeout', type=int, default=300, help='超过多少时间后，不再继续等待命令结果')
    parser.add_argument('--pool', type=int, default=50, help='同时进行的任务的数量')
    parser.add_argument('--show_query', action='store_true', help='只打印KQL查询语句，检查查询语句问题')
    parser.add_argument('--limit', type=int, dest='dev_limit', help='调试用，过滤前几条记录执行任务')
    parser.add_argument('--show_count', action='store_true', help='每行记录去打印序号')
    parser.add_argument('--csvfile', dest='dev_csvfile', help="仅限本地使用，输出一份到文件")

    args = parser.parse_args(args)
    return args

def main(args: Optional[List[str]] = None) -> None:
    """主函数"""
    # 初始化命令行参数
    if args is None:
        # 蓝鲸作业平台有BUG，重做的任务，空着的参数里都会被填充null
        args = [arg for arg in sys.argv[1:] if arg != 'null' and not arg.endswith('=null')]
    args = init_argparse(args)

    envs = {}
    if 'JENKINS_HOME' not in os.environ:
        if not args.envfile:
            return job_fail("请提供存有用户认证信息的env文件的文件名")
        # 初始化环境变量
        envs = init_environs(args.envfile)
        if not envs:
            return job_fail("env文件解密错误")
        if envs['runtime_env'] == 'blueking':
            for key in args.__dict__.keys():
                if key.startswith('dev_'):
                    v = getattr(args, key)
                    if v:
                        print(f'{key}={v}: 不支持的命令行参数，已经忽略。', file=sys.stderr)
                        setattr(args, key, None)

    # 实例化客户端
    CLOUD = getattr(msrestazure.azure_cloud, os.environ.get('CLOUD_NAME', envs.get('cloud_name')))
    client = Azure_Client(CLOUD)

    if args.subscription:
        # pylint: disable = no-member
        subscription = client.get_subscriptionid_by_name(args.subscription)
        if not subscription:
            return job_fail("订阅或订阅ID不存在:%s" % args.subscription)
        args.subscription = subscription
    is_linux = SCRIPTS[args.command]['is_linux']
    if is_linux:
        script_lines = SCRIPTS[args.command]['cmd'] + ['check %s' % i for i in args.services]
    else:
        script_lines = SCRIPTS[args.command]['cmd'] + ['check("%s")' % i for i in args.services]
    base_query = str_query_base(is_linux, args.subscription, args.resource_group, args.filters)
    if args.show_query:
        print(base_query)  # 蓝鲸平台上，那么多内容放到下面的job_success里会返回错误。单独print出来。
        return job_success("已打印查询语句，程序结束。")

    # 检查虚拟机数量
    vm_count = get_vm_count(client, base_query)
    print("虚拟机数量:", vm_count)
    if not vm_count:
        return job_fail("未匹配到任何虚拟机，建议检查过滤语句。")

    # 查询虚拟机信息
    vm_info = get_vm_info(client, base_query, args.dev_limit)
    # UTC转为本地时间
    data_format.parse_time(vm_info.data)

    # 异步执行任务
    job = RuncmdAsyncJob(client, args.pool, args.timeout, show_count=args.show_count, csvfile=args.dev_csvfile)
    for vm in vm_info.data:
        # vm 是个字典，下面是全部的字段
        # vm_name = vm.get('vm_name')
        # vm_resource_group = vm.get('resourceGroup')
        # vm_ostype = vm.get('OSType')
        # vm_subscription_id = vm.get('subscriptionId')
        # vm_subscription_name = vm.get('SubName')
        # vm_ip_address = vm.get('ip_address')
        # vm_osdisk_creattime = vm.get('osdisk_creattime')
        pool_is_full = job.add_task(vm, args.services, script_lines)
        if pool_is_full:
            job.get_task_result()
    # 最后取完所有的任务结果
    while job.get_task_result():
        pass
    job.close()


SCRIPTS = {
    'is-enabled': {'is_linux': True, 'cmd': textwrap.dedent("""\
        function check() {
            res=$($cmd $1)
            [ "$res" != "" ] && echo "$res" || echo "ReturnNone"
        }
        systemctl --version &>/dev/null
        if [ $? -eq 0 ]; then
            cmd="systemctl is-enabled"
        else
            cmd="chkconfig --list"
        fi

        """).splitlines()},
    # 只要进程名和状态，隐藏第一行标题，精确匹配进程名和状态
    # ps -eo ucmd,s h | grep -e '^sshd\s\+[RSD]$'
    # pylint: disable = anomalous-backslash-in-string
    'ps': {'is_linux': True, 'cmd': textwrap.dedent("""\
        function check() {
            ps -eo ucmd,s h | grep -e '^'$1'\s\+[RSD]$' | wc -l
        }

        """).splitlines()},
    # 检查迅销Symantec，-p参数显示版本信息，-a参数显示状态Enabled
    # 因为参数前面的横杠会先被命令行参数处理，导致参数解析错误，需要再包上一层引号
    'symantec_info': {'is_linux': True, 'cmd': textwrap.dedent("""\
        function check() {
            res=$($cmd $1)
            [ "$res" != "" ] && echo "$res" || echo "ReturnNone"
        }
        cmd="/opt/Symantec/symantec_antivirus/sav info"
                                                               
        """).splitlines()},
    # 迅销统计Linux虚拟机omsagent安装的版本
    # 管道符不能放到命令里，放进去的话，就会当做前面命令的位置参数
    'rpm_qa': {'is_linux': True, 'cmd': textwrap.dedent("""\
        function check() {
            res=$($cmd)
            ([ "$res" != "" ] && echo "$res" | grep $1 || echo "ReturnNone") | head -n 1
        }
        cmd="rpm -qa"
                                                        
        """).splitlines()},
    # eval 任意执行命令的方式，由于可以执行危险命令，不可控，弃用。
    'Get-Process': {'is_linux': False, 'cmd': textwrap.dedent("""\
        function check($process) {
            $result = Get-Process $process -ErrorAction SilentlyContinue
            if ($?) {
                Write-Host Running
            } else {
                Write-Host NotRun
            }
        }

        """).splitlines()},
    'Get-Service': {'is_linux': False, 'cmd': textwrap.dedent("""\
        function check($service) {
            $result = Get-Service $service -ErrorAction SilentlyContinue
            if ($?) {
                $result | %{$_.Status}
            } else {
                Write-Host NotExist
            }
        }

        """).splitlines()},
    'Get-WindowsFeature': {'is_linux': False, 'cmd': textwrap.dedent("""\
        function check($service) {
            $result = Get-WindowsFeature -Name $service -ErrorAction SilentlyContinue
            if ($?) {
                $result | %{$_.InstallState}
            } else {
                Write-Host NotExist
            }
        }

        """).splitlines()},
    # 检查 XZ-Utils的版本
    # 这个命令没有参数，但是返回2行，在命令行参数分别填 XZ-Utils 和 liblzma，结果打印会比较清晰
    # 没有考虑返回不是2行的情况
    'xz-V': {'is_linux': True, 'cmd': textwrap.dedent("""\
        function check() {
            xz -V
        }
                                                      
        """).splitlines()},
    # 检查ssh的版本
    # 这个命令没有参数，在命令行参数填随便填一个，比如 version
    'ssh-V': {'is_linux': True, 'cmd': textwrap.dedent("""\
        function check() {
            ssh -V
        }
                                                      
        """).splitlines()},
    # 创建用户 Linux
    'useradd': {'is_linux': True, 'cmd': textwrap.dedent("""\
        check() {
            username=$(echo $1 | cut -d':' -f1)
            password=$(echo $1 | cut -d':' -f2-)
            useradd -m -p $(openssl passwd -1 "$password") $username
            if [ $? -eq 0 ]; then
                echo Done
            else
                echo Failed
            fi
        }
                                                         
        """).splitlines()},
    # 创建用户 Windows
    'New-LocalUser': {'is_linux': False, 'cmd': textwrap.dedent("""\
        function check($userandpass) {
            $username, $password, $group = $userandpass -split ":", 3
            New-LocalUser -Name "$username" -Password ("$password" | ConvertTo-SecureString -AsPlainText -Force)
            if ($group) {
                Add-LocalGroupMember -Group "$group" -Member "$username"
            }
            if ($?) {
                Write-Host Done
            } else {
                Write-Host Failed
            }
        }

        """).splitlines()},
    }


if __name__ == '__main__':

    job_start()

    ###### 可在此处开始编写您的脚本逻辑代码
    ###### iJobs中执行脚本成功和失败的标准只取决于脚本最后一条执行语句的返回值
    ###### 如果返回值为0，则认为此脚本执行成功，如果非0，则认为脚本执行失败

    main()

